#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import torch
import torchvision
import cv2
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim


# In[ ]:


my_alexnet = torchvision.models.alexnet(pretrained=True)


# In[ ]:


import torchvision.transforms as transforms

my_tf = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261])])
## 同学们可以针对训练时的transform做一些增广，比如随机旋转等


# In[ ]:


train_dataset = torchvision.datasets.CIFAR10(root='E:/ml/cifar',train=True,transform=my_tf,download=True)
test_dataset = torchvision.datasets.CIFAR10(root='E:/ml/cifar',train=False,transform=my_tf,download=True)


# In[ ]:


train_dataloader = DataLoader(dataset=train_dataset,batch_size=32,shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset)


# In[ ]:


for param in my_alexnet.parameters():
    param.requires_grad=False


# In[ ]:


my_alexnet


# In[ ]:


in_f = my_alexnet.classifier[6].in_features
my_alexnet.classifier[6] = nn.Linear(in_f,10)


# In[ ]:


learn_rate = 0.001
num_epoches = 1
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(my_alexnet.classifier[6].parameters(),lr=learn_rate,momentum=0.9)


# In[ ]:


device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')


# In[ ]:


#train
my_alexnet.to(device)
my_alexnet.train()
for epoch in range(num_epoches):
    print(f"epoch: {epoch+1}")
    for idx,(img,label)in enumerate(train_dataloader):
        images = img.to(device)
        labels = label.to(device)
        output = my_alexnet(images)
        loss = criterion(output,labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        if i%100==0:
            print(f"current loss = {loss.item()}")


# In[ ]:


#eval
my_alexnet.to(device)
my_alexnet.eval()
total,correct = 0 , 0
for images,label in test_dataloader:
    images = img.to(device)
    labels = label.to(device)
    output = my_alexnet(images)
    _,idx = torch.max(output.data,1)
    total += labels.size(0)
    correct +=(idx==labels).cpu().sum()

print(f"accuracy:{correct/total}")
      


# In[ ]:




